from fastapi import APIRouter, WebSocket
import uuid
import json
import time

ws_router = APIRouter()

class WebSocketService:
    def __init__(self):
        self.connections = {}

    async def connect(self, clientId: str, ws:WebSocket):
        await ws.accept()
        self.connections[clientId] = ws

    async def disconnect(self, clientId: str):
        ws = self.connections[clientId]
        await ws.close()
        del self.connections[clientId]

    async def broadcast_message(self, data: dict):
        for cid in self.connections:
            await self.connections[cid].send_json(data)

wsService = WebSocketService()

async def myjob():
    msg = {"message":"hello world! " + time.strftime("%Y-%m-%d, %H:%M:%S")}
    await wsService.broadcast_message(msg)

@ws_router.websocket(path="/ws")
async def ws_websocket(websocket:WebSocket):
    clientId = str(uuid.uuid4())
    await wsService.connect(clientId, websocket)

    try:
        while True:
            data = await websocket.receive_text()
            print("==============>", data)
            msg = json.loads(data)
            await wsService.broadcast_message(msg)
    except Exception as e:
        print("Error========>", str(e))
        await wsService.disconnect(clientId)
